package cn.cew.bytespy.transfer;

import cn.cew.bytespy.comm.AgentContextCache;
import cn.cew.bytespy.utils.BaseUtils;
import javassist.ClassPool;
import javassist.CtClass;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;

public class RecoverAgentTransformer implements ClassFileTransformer {

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {

        ClassPool classPool = ClassPool.getDefault();
        CtClass ctClass;
        String cName = BaseUtils.packageToDot(className);
        byte[] bytes = AgentContextCache.getOriginalBytes(cName);

        if (null == bytes) {
            return classfileBuffer;
        }
        try {
            ctClass = classPool.get(cName);
            ctClass.defrost();
            InputStream inputStream;
            inputStream = new ByteArrayInputStream(bytes);
            ctClass = classPool.makeClass(inputStream);
        } catch (Exception e) {
            e.printStackTrace();
            return classfileBuffer;
        }
        try {
            return ctClass.toBytecode();
        } catch (Exception e) {
            e.printStackTrace();
            System.out.println("ex at toBytecode");
            return classfileBuffer;
        }
    }
}
